import os.path

import pandas as pd
import torch
from numpy import ndarray
from torch import Tensor

from utils.file import get_root_path
from utils.vectors import get_ref_vectors
from websocket.msghandler.multipleExcelHandler import MultipleExcelHandler

file_path = get_root_path() + r'\test_dec_data'


def tensor2excel(data: Tensor | ndarray, file_name=None, header=False, index=None, index_label=None):
    # 将 PyTorch 张量转换为 NumPy 数组
    if isinstance(data, ndarray):
        pass
    elif isinstance(data, Tensor):
        data = data.numpy()
    else:
        raise ValueError("data type error,just support tensor or ndarray")
    # 创建一个数据框
    df = pd.DataFrame(data, index=index)
    abs_file_path = MultipleExcelHandler.get_file_path(file_path, file_name, suffix='.xlsx')
    # 保存数据框为 Excel 文件
    df.to_excel(abs_file_path, index=False, header=header, index_label=index_label)


def excel2tensor(file_name):
    # 从 Excel 文件中读取数据
    df = pd.read_excel(os.path.join(file_path, file_name), header=None)

    # 提取数据并转换为 NumPy 数组
    data_array = df.values

    # 将 NumPy 数组转换为 PyTorch 张量
    return torch.tensor(data_array, dtype=torch.double)


if __name__ == '__main__':
    tensor2excel(get_ref_vectors(1e4, 2))
